import numpy as np
import matplotlib.pyplot as plt

# Compute damping coefficients
def compute_damping_coefficients(n, omega, v1, v2):
    k = np.linspace(1, n, 1000)
    lambda_k = 1 - omega + omega * np.cos(k * np.pi / n)
    lambda_k_prime = 1 - omega + omega * np.cos((n - k) * np.pi / n)
    s_k = np.sin(k * np.pi / (2 * n))**2
    c_k = np.cos(k * np.pi / (2 * n))**2

    c1 = lambda_k**(v1 + v2) * s_k
    c2 = lambda_k**v1 * lambda_k_prime**v2 * s_k
    c3 = lambda_k_prime**v1 * lambda_k**v2 * c_k
    c4 = lambda_k_prime**(v1 + v2) * c_k

    return c1, c2, c3, c4

# Plot damping coefficients
def plot_damping_coefficients(n, omega):
    v_combinations = [[0, 0, 0], [0, 2, 1], [1, 1, 2], [2, 0, 3], [2, 2, 4], [4, 0, 5]]
    w = np.linspace(1, n//2, 1000)
    t = np.linspace(n//2, n, 1000)
    # 创建 3 行 2 列的子图
    fig, ax = plt.subplots(3, 2, figsize=(8,6))
    ax = ax.flatten()  # 将二维的轴数组展平为一维，方便索引

    for v1, v2, i in v_combinations:
        c1, c2, c3, c4 = compute_damping_coefficients(n, omega, v1, v2)
        ax[i].plot(t, c1, label=f'c1')
        ax[i].plot(t, c4, label=f'c4')
        ax[i].plot(n//2-w, c2, label=f'c2')
        ax[i].plot(n//2-w, c3, label=f'c3')
        ax[i].set_xlabel('Wavenumber k')
        ax[i].set_ylabel('Damping Coefficient')
        ax[i].set_title(f'Damping Coefficients for v1={v1}, v2={v2}')
        ax[i].legend()

    plt.tight_layout()  # 自动调整子图间距
    plt.savefig(f"./E9_41_{n}.jpg")

# Parameters
n = 64
omega = 2 / 3

# Plot damping coefficients for n=64
plot_damping_coefficients(n, omega)

# Repeat for n=128
n = 128
plot_damping_coefficients(n, omega)